package com.example.endusation.milvus;

import ai.djl.translate.TranslateException;
import com.mysql.cj.QueryResult;
import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.exception.IllegalResponseException;
import io.milvus.exception.ParamException;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.naming.directory.SearchResult;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import static java.util.Collections.singletonList;
@Service
public class MilvusService {

    private final MilvusClient client;

    private static final String COLLECTION_NAME = "discrete_math";
    private static final String VIDEO = "video";
    private static final String DOCUMENT = "document";
    private static final String VECTOR_FIELD_NAME = "embedding";

    public MilvusService(@Value("${milvus.host}") String host,
                         @Value("${milvus.port}") int port) {
        // 修正：使用 MilvusServiceClient，ConnectParam 用 newBuilder()
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost(host)
                .withPort(port)
                .build();
        this.client = new MilvusServiceClient(connectParam);
    }

    // 创建集合（定义字段结构，包含向量维度）
    public void createCollection(String collectionName) {
        // 1. 定义字段（必须包含主键和向量字段）
        List<FieldType> fields = new ArrayList<>();

        // 主键字段（Long 类型，自动生成）
        fields.add(FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build());

        // 文本内容字段（字符串）
        fields.add(FieldType.newBuilder()
                .withName("text")
                .withDataType(DataType.VarChar)
                .withMaxLength(2048) // 最大长度
                .build());

        // 向量字段（384 维，与 all-MiniLM-L6-v2 匹配）
        fields.add(FieldType.newBuilder()
                .withName(VECTOR_FIELD_NAME)
                .withDataType(DataType.FloatVector)
                .withDimension(384) // 修正：通过 FieldType 指定维度
                .build());

        // 源文件路径字段
        fields.add(FieldType.newBuilder()
                .withName("source_file")
                .withDataType(DataType.VarChar)
                .withMaxLength(1024)
                .build());

        // 位置信息（页码或时间戳）
        fields.add(FieldType.newBuilder()
                .withName("position")
                .withDataType(DataType.Int32)
                .build());

        // 2. 创建集合
        CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldTypes(fields) // 传入字段定义
                .withShardsNum(2) // 分片数量
                .build();

        R<RpcStatus> response = client.createCollection(createParam);
        if (response.getData()!=null) {
            System.out.println("集合创建成功！");
        } else {
            System.out.println("集合创建失败：" + response.getStatus());
        }
    }

    // 新增：加载集合到内存（搜索前必须执行）
    public void loadCollection(String collectionName) {




        // 步骤2：加载集合到内存（默认加载所有分区，也可指定分区）
        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)


                .build();

        R<RpcStatus> loadResult = client.loadCollection(loadParam);
        if (loadResult.getStatus() != 0) {
            throw new RuntimeException("加载集合 " + collectionName + " 到内存失败：" + loadResult.getMessage());
        }
        System.out.println("集合 " + collectionName + " 成功加载到内存");
    }
    // 插入向量数据
    public long insert(String text, float[] vector, String sourceFile, int pageOrTime,String collectionName) {
        // 1. 处理向量：先归一化，再转为 List<Float>（核心修正：新增归一化步骤）
        float[] normalizedVector = normalizeVector(vector); // 第一步：归一化原始 float[] 向量
        List<Float> vectorList = new ArrayList<>(normalizedVector.length);
        for (float value : normalizedVector) { // 基于归一化后的向量构建 List
            vectorList.add(value);
        }

        // 2. 按「字段维度」组织数据（这部分逻辑不变）
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("text", Collections.singletonList(text)));
        fields.add(new InsertParam.Field("embedding", Collections.singletonList(vectorList))); // 插入归一化后的向量
        fields.add(new InsertParam.Field("source_file", Collections.singletonList(sourceFile)));
        fields.add(new InsertParam.Field("position", Collections.singletonList(pageOrTime)));

        // 3. 构建 InsertParam 并执行插入（逻辑不变）
        try {
            InsertParam insertParam = InsertParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withFields(fields)
                    .build();

            R<MutationResult> response = client.insert(insertParam);
            if (response.getStatus() == 0) {
                return response.getData().getInsertCnt(); // 注意：getInsertCnt() 返回插入数量，若需ID需用 getInsertIdsList()
            } else {
                throw new RuntimeException("Milvus 插入失败：状态码=" + response.getStatus() + "，原因=" + response.getMessage());
            }
        } catch (ParamException e) {
            throw new RuntimeException("插入参数错误：" + e.getMessage(), e);
        }
    }

    /**
     * 核心工具方法：对 float[] 类型向量进行 L2 范数归一化（转为单位向量）
     * @param originalVector 原始未归一化的向量（float[]）
     * @return 归一化后的单位向量（float[]）
     */
    private float[] normalizeVector(float[] originalVector) {
        // 校验输入：避免空向量或null
        if (originalVector == null || originalVector.length == 0) {
            throw new IllegalArgumentException("原始向量不能为 null 或空数组");
        }

        // 步骤1：计算原始向量的 L2 范数（平方和的平方根）
        double sumOfSquares = 0.0;
        for (float value : originalVector) {
            sumOfSquares += value * value; // 累加每个元素的平方
        }
        double l2Norm = Math.sqrt(sumOfSquares); // 开平方得到 L2 范数

        // 步骤2：处理特殊情况（若向量全为0，范数接近0，直接返回原向量避免除以0）
        if (l2Norm < 1e-6) { // 1e-6 是极小值，避免浮点计算误差导致的除以0
            return originalVector.clone(); // 克隆避免修改原始数组
        }

        // 步骤3：归一化：每个元素除以 L2 范数，得到单位向量
        float[] normalizedVector = new float[originalVector.length];
        for (int i = 0; i < originalVector.length; i++) {
            normalizedVector[i] = (float) (originalVector[i] / l2Norm);
        }

        return normalizedVector;
    }

    // 搜索相似文本
    public List<QueryResult> search(String queryText, EmbeddingService embeddingService, int topK,String collectionName)
            throws TranslateException {
        loadCollection(collectionName);
        // 1. 生成查询向量 + 关键修正：添加向量归一化（与存储时逻辑一致）
        // 校验输入参数，避免空指针
        if (queryText == null || queryText.trim().isEmpty()) {
            throw new IllegalArgumentException("查询文本（queryText）不能为null或空字符串");
        }
        if (embeddingService == null) {
            throw new IllegalArgumentException("EmbeddingService不能为null");
        }
        if (topK <= 0) {
            throw new IllegalArgumentException("TopK必须为正整数（当前值：" + topK + "）");
        }

        // 1.1 调用Embedding服务生成原始向量
        float[] originalQueryVector = embeddingService.getEmbedding(queryText);
        if (originalQueryVector == null || originalQueryVector.length == 0) {
            throw new RuntimeException("Embedding服务生成的查询向量为空，请检查服务可用性");
        }

        // 1.2 对原始向量进行L2归一化（核心：确保与存储的向量归一化逻辑一致，IP才能等价余弦相似度）
        float[] normalizedQueryVector = normalizeVector(originalQueryVector);

        // 1.3 归一化向量转为List<Float>（适配Milvus的withVectors参数要求）
        List<Float> queryVectorList = new ArrayList<>(normalizedQueryVector.length);
        for (float value : normalizedQueryVector) {
            queryVectorList.add(value);
        }

        // 2. 构建搜索参数 + 关键修正：withOutFields参数改为List<String>格式

            SearchParam searchParam = SearchParam.newBuilder()
                    .withCollectionName(collectionName) // 确保COLLECTION_NAME是已定义的常量（非null/空）
                    .withVectorFieldName(VECTOR_FIELD_NAME) // 确保与集合中向量字段名一致（如"embedding"）
                    .withVectors(Collections.singletonList(queryVectorList)) // 传入归一化后的查询向量列表
                    .withTopK(topK) // 已校验为正整数，避免ParamException
                    .withMetricType(MetricType.IP) // IP内积（等价余弦相似度，依赖归一化）
                    .withOutFields(Arrays.asList("text", "source_file", "position")) // 修正：参数转为List<String>
                    .build();

        // 3. 执行搜索
        R<SearchResults> response = client.search(searchParam);
        if (response.getStatus()!=0) {
            throw new RuntimeException("搜索失败：" + response.getStatus());
        }
        // 4. 解析结果
        // 解析结果
        SearchResults searchResults = response.getData(); // 外层对象
        SearchResultData resultData = searchResults.getResults();
        SearchResultsWrapper wrapper = new SearchResultsWrapper(resultData);
        int indexOfTarget = 0; // 单向量查询，固定为0

        try {
            List<SearchResultsWrapper.IDScore> idScores = wrapper.getIDScore(indexOfTarget);
            List<?> textList = wrapper.getFieldData("text", indexOfTarget);
            List<?> sourceFileList = wrapper.getFieldData("source_file", indexOfTarget);
            List<?> positionList = wrapper.getFieldData("position", indexOfTarget);

            List<QueryResult> finalResults = new ArrayList<>();
            for (int i = 0; i < idScores.size(); i++) {
                float score = idScores.get(i).getScore();
                long id = idScores.get(i).getLongID();
                String text = (String) textList.get(i);
                String sourceFile = (String) sourceFileList.get(i);
                int position = (Integer) positionList.get(i);

                finalResults.add(new QueryResult(id, text, sourceFile, position, score));
            }
            return finalResults;

        } catch (ParamException | IllegalResponseException e) {
            throw new RuntimeException("解析搜索结果失败：" + e.getMessage(), e);
        }
    }

    // 自定义查询结果类（根据需要调整）
    public static class QueryResult {
        private long id;           // 新增：向量ID（自增ID，若为字符串则改为 String id）
        private String text;       // 文本
        private String sourceFile; // 源文件
        private int position;      // 位置
        private float score;       // 相似度分数

        // 修改后的构造方法：接受5个参数
        public QueryResult(long id, String text, String sourceFile, int position, float score) {
            this.id = id;
            this.text = text;
            this.sourceFile = sourceFile;
            this.position = position;
            this.score = score;
        }

        // 新增 getter 方法（按需添加）
        public long getId() {
            return id;
        }


        // Getter 方法
        public String getText() { return text; }
        public String getSourceFile() { return sourceFile; }
        public int getPosition() { return position; }
        public float getScore() { return score; }
    }
}
